#include "StdAfx.h"
#include "StaticHookDll.h"
#include "detours.h"

CHAR StaticHookDll::dllPath[MAX_PATH] = "taxer.dll";
static BOOL CALLBACK ExportCallback(PVOID pContext,
									ULONG nOrdinal,
									PCHAR pszSymbol,
									PVOID pbTarget)
{
	(void)pContext;
	(void)pbTarget;
	(void)pszSymbol;

	if (nOrdinal == 1) {
		*((BOOL *)pContext) = TRUE;
	}
	return TRUE;
}

BOOL StaticHookDll::DoesDllExportOrdinal1(PCHAR pszDllPath)
{
	HMODULE hDll = LoadLibraryEx(pszDllPath, NULL, DONT_RESOLVE_DLL_REFERENCES);
	if (hDll == NULL) {
		/*printf("withdll.exe: LoadLibraryEx(%s) failed with error %d.\n",
			pszDllPath,
			GetLastError());*/
		return FALSE;
	}

	BOOL validFlag = FALSE;
	DetourEnumerateExports(hDll, &validFlag, ExportCallback);
	FreeLibrary(hDll);
	return validFlag;
}

//////////////////////////////////////////////////////////////////////////////
//
static BOOL CALLBACK ListBywayCallback(PVOID pContext,
									   PCHAR pszFile,
									   PCHAR *ppszOutFile)
{
	(void)pContext;

	*ppszOutFile = pszFile;
	/*if (pszFile) {
		printf("    %s\n", pszFile);
	}*/
	return TRUE;
}

static BOOL CALLBACK ListFileCallback(PVOID pContext,
									  PCHAR pszOrigFile,
									  PCHAR pszFile,
									  PCHAR *ppszOutFile)
{
	(void)pContext;

	*ppszOutFile = pszFile;
	//printf("    %s -> %s\n", pszOrigFile, pszFile);
	return TRUE;
}

static BOOL CALLBACK AddBywayCallback(PVOID pContext,
									  PCHAR pszFile,
									  PCHAR *ppszOutFile)
{
	PBOOL pbAddedDll = (PBOOL)pContext;
	if (!pszFile && !*pbAddedDll) {                     // Add new byway.
		*pbAddedDll = TRUE;
		*ppszOutFile = StaticHookDll::dllPath;
	}
	return TRUE;
}

FAXUSBMANAGER_API bool ResetDll( char *appPathName )
{
	StaticHookDll hooker;
	BOOL ret = hooker.SetDll(appPathName, TRUE);
	if (!ret)
	{
		return false;
	}
	return true;
}
StaticHookDll::StaticHookDll(void)
{
	//m_isRemove = FALSE;
}

StaticHookDll::~StaticHookDll(void)
{
}

BOOL StaticHookDll::SetDll( char *appPathName, BOOLEAN remove )
{
	BOOL bGood = TRUE;
	HANDLE hOld = INVALID_HANDLE_VALUE;
	HANDLE hNew = INVALID_HANDLE_VALUE;
	PDETOUR_BINARY pBinary = NULL;
	//m_isRemove = remove;
	CHAR szOrg[MAX_PATH];
	CHAR szNew[MAX_PATH];
	CHAR szOld[MAX_PATH];

	szOld[0] = '\0';
	szNew[0] = '\0';

#ifdef _CRT_INSECURE_DEPRECATE
	strcpy_s(szOrg, sizeof(szOrg), appPathName);
	strcpy_s(szNew, sizeof(szNew), szOrg);
	strcat_s(szNew, sizeof(szNew), "#");
	strcpy_s(szOld, sizeof(szOld), szOrg);
	strcat_s(szOld, sizeof(szOld), "~");
#else
	strcpy(szOrg, pszPath);
	strcpy(szNew, szOrg);
	strcat(szNew, "#");
	strcpy(szOld, szOrg);
	strcat(szOld, "~");
#endif
	//printf("  %s:\n", pszPath);

	hOld = CreateFile(szOrg,
		GENERIC_READ,
		FILE_SHARE_READ,
		NULL,
		OPEN_EXISTING,
		FILE_ATTRIBUTE_NORMAL,
		NULL);

	if (hOld == INVALID_HANDLE_VALUE) {
		/*printf("Couldn't open input file: %s, error: %d\n",
			szOrg, GetLastError());*/
		bGood = FALSE;
		goto end;
	}

	hNew = CreateFile(szNew,
		GENERIC_WRITE | GENERIC_READ, 0, NULL, CREATE_ALWAYS,
		FILE_ATTRIBUTE_NORMAL | FILE_FLAG_SEQUENTIAL_SCAN, NULL);
	if (hNew == INVALID_HANDLE_VALUE) {
		/*printf("Couldn't open output file: %s, error: %d\n",
			szNew, GetLastError());*/
		bGood = FALSE;
		goto end;
	}

	if ((pBinary = DetourBinaryOpen(hOld)) == NULL) {
		//printf("DetourBinaryOpen failed: %d\n", GetLastError());
		goto end;
	}

	if (hOld != INVALID_HANDLE_VALUE) {
		CloseHandle(hOld);
		hOld = INVALID_HANDLE_VALUE;
	}

	{
		BOOL bAddedDll = FALSE;

		DetourBinaryResetImports(pBinary);

		if (!remove) {
			if (!DetourBinaryEditImports(pBinary,
				&bAddedDll,
				AddBywayCallback, NULL, NULL, NULL)) {
					DWORD err = GetLastError();
					//printf("DetourBinaryEditImports failed: %d\n", GetLastError());
			}
		}

		if (!DetourBinaryEditImports(pBinary, NULL,
			ListBywayCallback, ListFileCallback,
			NULL, NULL)) {
				DWORD err = GetLastError();
				//printf("DetourBinaryEditImports failed: %d\n", GetLastError());
		}

		if (!DetourBinaryWrite(pBinary, hNew)) {
			DWORD err = GetLastError();
			//printf("DetourBinaryWrite failed: %d\n", GetLastError());
			bGood = FALSE;
		}

		DetourBinaryClose(pBinary);
		pBinary = NULL;

		if (hNew != INVALID_HANDLE_VALUE) {
			CloseHandle(hNew);
			hNew = INVALID_HANDLE_VALUE;
		}

		if (bGood) {
			if (!DeleteFile(szOld)) {
				DWORD dwError = GetLastError();
				if (dwError != ERROR_FILE_NOT_FOUND) {
					//printf("Warning: Couldn't delete %s: %d\n", szOld, dwError);
					bGood = FALSE;
				}
			}
			if (!MoveFile(szOrg, szOld)) {
				/*printf("Error: Couldn't back up %s to %s: %d\n",
					szOrg, szOld, GetLastError());*/
				bGood = FALSE;
			}
			if (!MoveFile(szNew, szOrg)) {
				/*printf("Error: Couldn't install %s as %s: %d\n",
					szNew, szOrg, GetLastError());*/
				bGood = FALSE;
			}
		}

		DeleteFile(szNew);
		DeleteFile(szOld);
	}


end:
	if (pBinary) {
		DetourBinaryClose(pBinary);
		pBinary = NULL;
	}
	if (hNew != INVALID_HANDLE_VALUE) {
		CloseHandle(hNew);
		hNew = INVALID_HANDLE_VALUE;
	}
	if (hOld != INVALID_HANDLE_VALUE) {
		CloseHandle(hOld);
		hOld = INVALID_HANDLE_VALUE;
	}
	return bGood;
}